import os

from langchain.chains.sql_database.query import create_sql_query_chain
from langchain_community.agent_toolkits import SQLDatabaseToolkit
from langchain_community.tools import QuerySQLDataBaseTool
from langchain_community.utilities import SQLDatabase
from langchain_core.messages import SystemMessage, HumanMessage
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_openai import ChatOpenAI
from langgraph.prebuilt import chat_agent_executor

# sqlalchemy连数据库
from venv.Lib.operator import itemgetter

HOSTNAME = '127.0.0.1'
PORT = '3306'
DATABASE = '******'
USERNAME = '******'
PASSWORD = '******'
MYSQL_URL = 'mysql+mysqldb://{}:{}@{}:{}/{}?charset=utf8mb4'.format()

db = SQLDatabase.from_uri(MYSQL_URL)

# 测试连接是否用用
print(db.get_usable_table_names())

print(db.run('select * from database_source_name limit 5;'))

os.environ['LANGCHAIN_TRACING_V2'] = "true"
os.environ['LANGCHAIN_API_KEY'] = '1123'

# 调用大模型

model = ChatOpenAI(model='gpt-4-turbo')

# test_chain = create_sql_query_chain(model, db)
#
# resp = test_chain.invoke({'question':'请问员工表中有多少条数据？'})
#
# print(resp)
#
# answer_prompt = PromptTemplate.from_template(
#     """给定以下用户问题、SQL语句和SQL执行后的结果，回答用户问题，
#     Question: {question}
#     SQL Query: {query}
#     SQL Result: {result}
#     回答："""
# )
#
# execute_sql_test = QuerySQLDataBaseTool(db=db)
# #断言
# chain = (RunnablePassthrough.assign(query=test_chain).assign(result=itemgetter('query') | execute_sql_test) | answer_prompt | model | StrOutputParser())
#
# rep = chain.invoke(input={'question':'请问员工表中有多少条数据?'})
#
# print(rep)


# 用agent实现查询数据库sql并执行
db1 = SQLDatabase.from_uri(MYSQL_URL)

toolkit = SQLDatabaseToolkit(db=db1, llm=model)
tools = toolkit.get_tools()

system_prompt = """
您是一个被设计用来与SQL数据库交互的代理。
约定一个输入问题，创建一个语法正确的SQL语句并执行，然后查询结果并返回答案。
除非用于制定了他们想要获得的示例的具体数量，否则始终将SQL查询限制为最多10个结果
你可以对相关结果进行排序，以返回MySQL数据库中最匹配的数据。
您可以使用与数据库交互的工具，在执行查询之前，你必须仔细检查，如果在执行查询时出现错误，请重写查询并重试
不要对数据库做任何DML语句

首先，你要查看数据库中的表，看看可以查询什么
不要跳过这一步
然后查询相关的表的模式
"""
system_message = SystemMessage(content=system_prompt)

agent_executor = chat_agent_executor.create_tool_calling_executor(model, tools, system_message)

resp = agent_executor.invoke({'messages': [HumanMessage(content='请问有多少员工？')]})

result = resp['messages']

print(result)
print(result[len(result) - 1])
